﻿using StableDiffusionTools.Extensions;
using System;
using System.Collections.Generic;
using System.Drawing;
using System.IO;
using System.Linq;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Threading.Tasks;

namespace StableDiffusionTools.Framework
{
    public class Safetensors
    {
        public Bitmap[] CoverImages = null;
        public string Metadata { get; private set; }
        public string MetaModel { get; private set; }
        public int MaxTagLength { get; private set; } = 0;
        public Result CoverResult { get; private set; } = Result.Missing;

        private bool parseTagsDictionary = true;
        private Dictionary<string, Tuple<int, int>> tagsDictionary = null;
        public Dictionary<string, Tuple<int, int>> TagsDictionary
        {
            get
            {
                if (tagsDictionary == null && parseTagsDictionary)
                {
                    parseTagsDictionary = false;
                    if (!string.IsNullOrEmpty(Metadata))
                    {
                        if (TryExtractTags(Metadata, out Dictionary<string, Tuple<int, int>> t, out int maxLen))
                        {
                            tagsDictionary = t;
                            MaxTagLength = maxLen;
                        }
                    }
                }

                return tagsDictionary;
            }
        }

        private bool parseTags = true;
        private List<string> tags = null;
        public List<string> Tags
        {
            get
            {
                if (tags == null && parseTags)
                {
                    parseTags = false;
                    var tagsDict = TagsDictionary;
                    if(tagsDict != null)
                        tags = FormatTags(tagsDict, MaxTagLength);
                }

                return tags;
            }
        }

        public List<Tuple<string, string>> Tensors { get; set; }

        public Safetensors(string meta, string model, Bitmap[] coverImages, List<Tuple<string, string>> tensors, Result result)
        {
            Metadata = meta;
            MetaModel = model;
            CoverImages = coverImages;
            Tensors = tensors ?? new List<Tuple<string, string>>();
            CoverResult = result;
        }

        public static Safetensors LoadSafetensors(string filePath, bool loadTensors = false)
        {
            using (FileStream file = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read))
            {
                using (var sr = new BinaryReader(file))
                {
                    Result result = Result.Missing;
                    ulong headerSize = sr.ReadUInt64();
                    byte[] header = sr.ReadBytes((int)headerSize);

                    string metaStr = "";
                    string metaModel = "";
                    Bitmap[] coverImages = null;
                    List<Tuple<string, string>> tensors = new List<Tuple<string, string>>();

                    Utf8JsonReader reader = new Utf8JsonReader(header);
                    if(JsonDocument.TryParseValue(ref reader, out JsonDocument doc))
                    {
                        if (loadTensors)
                        {
                            try
                            {
                                foreach (var item in doc.RootElement.EnumerateObject())
                                {
                                    if (!item.Name.Equals("__metadata__"))
                                    {
                                        TensorType dType = null;
                                        try
                                        {
                                            dType = JsonSerializer.Deserialize<TensorType>(item.Value.GetRawText());
                                            tensors.Add(new Tuple<string, string>(item.Name, dType.DType));
                                            //System.Diagnostics.Debug.WriteLine($"{item.Name} {dType.DType}");
                                        }
                                        catch (Exception) { }
                                    }
                                }
                            }
                            catch (Exception) { }
                        }

                        if (doc.RootElement.TryGetProperty("__metadata__", out JsonElement metadata))
                        {
                            JsonObject jo = JsonObject.Create(metadata);
                            if (jo.TryGetPropertyValue("ssmd_cover_images", out JsonNode node))
                            {
                                string nodeStr = node.ToString();
                                try
                                {
                                    if (nodeStr.Length > 4)
                                    {
                                        JsonArray jArr = (JsonArray)JsonNode.Parse(nodeStr);
                                        coverImages = new Bitmap[jArr.Count];
                                        MemoryStream[] mem = new MemoryStream[jArr.Count];
                                        Image[] img = new Image[jArr.Count];
                                        try
                                        {
                                            for (int i = 0; i < jArr.Count; i++)
                                            {
                                                string coverStr = jArr[i].ToString();
                                                var arr = Convert.FromBase64String(coverStr);
                                                mem[i] = new MemoryStream(arr, false);
                                                img[i] = Image.FromStream(mem[i]);
                                                coverImages[i] = ((Bitmap)img[i]).Clone(new Rectangle(0, 0, img[i].Width, img[i].Height), System.Drawing.Imaging.PixelFormat.Format32bppPArgb);
                                                arr = null;
                                            }
                                            if (jArr.Count > 0)
                                                result = Result.Found;
                                            else
                                                result = Result.Empty;
                                        }
                                        finally
                                        {
                                            for (int i = 0; i < jArr.Count; i++)
                                            {
                                                mem[i]?.Dispose();
                                                img[i]?.Dispose();
                                            }
                                        }
                                    }
                                    else
                                    {
                                        result = Result.Empty;
                                    }
                                }
                                catch (Exception)
                                {
                                    result = Result.Failed;
                                }
                                finally
                                {
                                    jo["ssmd_cover_images"] = result.ToString().ToUpper();
                                }
                            }

                            if (jo.TryGetPropertyValue("ss_sd_model_name", out JsonNode jnModelName))
                            {
                                metaModel = jnModelName.ToString();
                            }

                            metaStr = jo.ToStringIndented();
                        }
                        doc.Dispose();
                    }

                    return new Safetensors(metaStr, metaModel, result == Result.Found ? coverImages : null, tensors, result);
                }
            }
        }

        private static bool TryExtractTags(string meta, out Dictionary<string, Tuple<int, int>> tags, out int maxLen)
        {
            Dictionary<string, Tuple<int, int>> tagList = new Dictionary<string, Tuple<int, int>>();
            int overallMaxLen = 0;

            JsonElement ss_tags = new JsonElement();
            try
            {
                using (JsonDocument doc = JsonDocument.Parse(meta))
                {
                    if(doc.RootElement.TryGetProperty("ss_tag_frequency", out JsonElement elem)){
                        ss_tags = JsonDocument.Parse(elem.ToString()).RootElement;
                    }
                }
            }
            catch (Exception ex)
            {
            }

            if (ss_tags.ValueKind != JsonValueKind.Null && ss_tags.ValueKind != JsonValueKind.Undefined)
            {
                foreach (var item in ss_tags.EnumerateObject())
                {
                    int mul = 1;
                    if (item.Name.Contains('_') && int.TryParse(item.Name.Substring(0, item.Name.IndexOf('_')), System.Globalization.NumberStyles.None, System.Globalization.CultureInfo.InvariantCulture, out int mulOut))
                    {
                        mul = mulOut;
                    }

                    if (item.Value.ValueKind == JsonValueKind.Object)
                    {
                        foreach (var item2 in item.Value.EnumerateObject())
                        {
                            if (item2.Value.ValueKind == JsonValueKind.Number)
                            {
                                string tagName = item2.Name.Trim();
                                int num = item2.Value.GetInt32();
                                int numRepeats = mul * num;

                                if (tagList.ContainsKey(tagName))
                                {
                                    tagList[tagName] = new Tuple<int, int>(tagList[tagName].Item1 + num, tagList[tagName].Item2 + numRepeats);
                                }
                                else
                                {
                                    tagList.Add(tagName, new Tuple<int, int>(num, numRepeats));
                                }

                                if (overallMaxLen < tagName.Length)
                                    overallMaxLen = tagName.Length;
                            }
                        }
                    }
                }
            }

            tags = tagList;
            maxLen = overallMaxLen;
            return tagList.Count > 0;
        }

        public static List<string> FormatTags(Dictionary<string, Tuple<int, int>> tagList, int overallMaxLen)
        {
            List<string> result = new List<string>();
            if (tagList.Count > 0)
            {
                var sortedDict = from entry in tagList orderby entry.Value descending select entry;
                result.AddRange(sortedDict.Select(x => $"{x.Key.PadRight(overallMaxLen, ' ')} {x.Value}"));
            }
            return result;
        }

        public enum Result
        {
            Missing,
            Found,
            Failed,
            Empty
        }
    }
}
